package token_bucket

import (
	"context"
	_ "embed"
	"fmt"
	"ratelimit"
	"sync"
	"sync/atomic"
	"time"

	"github.com/go-redis/redis/v8"
	xrate "golang.org/x/time/rate"
)

const (
	tokenFormat     = "ratelimit.%s.tokens"
	timestampFormat = "ratelimit.%s.ts"
	pingInterval    = time.Millisecond * 100
)

//go:embed redis_token_bucket.lua
var redisTokenBucketLua string

type RedisTokenBucketLimiter struct {
	client          *redis.Client
	capacity        int
	rate            int // 单位秒
	tokenKey        string
	timestampKey    string
	redisAlive      uint32
	lock            sync.Mutex
	monitorStarted  bool
	fallbackLimiter *xrate.Limiter
}

func NewRedisTokenBucketLimiter(rate, capacity int, client *redis.Client, key string) ratelimit.Limiter {
	limiter := &RedisTokenBucketLimiter{
		client:          client,
		capacity:        capacity,
		rate:            rate,
		tokenKey:        fmt.Sprintf(tokenFormat, key),
		timestampKey:    fmt.Sprintf(timestampFormat, key),
		redisAlive:      1,
		fallbackLimiter: xrate.NewLimiter(xrate.Every(time.Second/time.Duration(rate)), capacity),
	}

	return limiter
}

func (r *RedisTokenBucketLimiter) Allow() bool {
	return r.allowN(1)
}

func (r *RedisTokenBucketLimiter) allowN(n int) bool {
	if atomic.LoadUint32(&r.redisAlive) == 0 {
		return r.fallbackLimiter.AllowN(time.Now(), n)
	}

	now := time.Now().UnixMilli()
	ctx := context.Background()
	cmd := r.client.Eval(ctx, redisTokenBucketLua, []string{r.tokenKey, r.timestampKey}, r.rate/1e3, r.capacity, now, n)
	result, err := cmd.Result()
	if err != nil {
		go r.startMonitor()

		return r.fallbackLimiter.AllowN(time.Now(), n)
	}

	ret, ok := result.(bool)
	if !ok {
		go r.startMonitor()
		return r.fallbackLimiter.AllowN(time.Now(), n)
	}
	return ret
}

/*
对于存在更复杂逻辑或者多个变量之间存在关系的场景（例如，你需要保证某几个变量的变化是一致的），原子操作往往不能替代锁。
在这段代码中，虽然monitorStarted本身可以使用原子操作来管理，但是你还需要确保在monitorStarted标志位设置之后，监视器确实启动了。
由于涉及到多步骤操作的同步，因此仍然使用锁可能是更安全、更直观的选择
*/
func (r *RedisTokenBucketLimiter) startMonitor() {
	r.lock.Lock()
	defer r.lock.Unlock()

	if r.monitorStarted {
		return
	}

	r.monitorStarted = true
	go r.keepAlive()
}

func (r *RedisTokenBucketLimiter) keepAlive() {
	ticker := time.NewTicker(pingInterval)

	defer func() {
		ticker.Stop()
		atomic.StoreUint32(&r.redisAlive, 1)
		r.lock.Lock()
		r.monitorStarted = false
		r.lock.Unlock()
	}()

	for range ticker.C {
		res := r.client.Ping(context.Background())
		if res.Err() == nil {
			return
		}
	}
}
